import numpy as np

class ETPModel:
    def __init__(self, C_a, C_m, H_m, U_a, COP, T_o):
        self.C_a = C_a
        self.C_m = C_m
        self.H_m = H_m
        self.U_a = U_a
        self.COP = COP
        self.T_o = T_o
        self.T_a = 25.0  # 初始室温
        self.T_m = 25.0  # 初始固体温度

    def step(self, Q_ac, Q_i, Q_s, dt):
        Q_a = Q_ac + 0.5 * Q_i + 0.5 * Q_s
        Q_m = 0.5 * Q_i + 0.5 * Q_s

        dTa_dt = (self.T_m * self.H_m - (self.U_a + self.H_m) * self.T_a + Q_a + self.T_o * self.U_a) / self.C_a
        dTm_dt = (self.H_m * (self.T_a - self.T_m) + Q_m) / self.C_m

        self.T_a += dTa_dt * dt
        self.T_m += dTm_dt * dt

        return self.T_a, self.T_m

    def compute_power(self, Q_ac):
        return Q_ac / self.COP

class SemiMarkovAC:
    def __init__(self, P_rate, lock_time, execution_period, etp_model, A=0.2, B=3.5):
        self.P_rate = P_rate
        self.lock_steps = int(lock_time / execution_period)
        self.etp = etp_model
        self.state = 'OFF'
        self.lock_timer = 0
        self.execution_period = execution_period

        # 用户功率期望函数参数
        self.A = A
        self.B = B

        self.power = 0.0

    def compute_expected_power(self):
        # todo:这个方法应该是模拟需求曲线, 输入虚拟价格, 输出动态功率.可以先做输出动态规律的部分
        # 公式: P_d(k) = A*T_a(k) + B*T_a(k-1) + C*T_m(k-1) + D(k)

        P_d = self.A * self.etp.T_a + self.B
        return min(max(P_d, 0.0), self.P_rate)

    def update_state(self, u_0, u_1):
        if self.state == 'ON':
            if np.random.rand() < u_1 * self.execution_period:
                self.state = 'OFF_LOCKED'
                self.lock_timer = self.lock_steps

        elif self.state == 'OFF_LOCKED':
            self.lock_timer -= 1
            if self.lock_timer <= 0:
                self.state = 'OFF'

        elif self.state == 'OFF':
            if np.random.rand() < u_0 * self.execution_period:
                self.state = 'ON_LOCKED'
                self.lock_timer = self.lock_steps

        elif self.state == 'ON_LOCKED':
            self.lock_timer -= 1
            if self.lock_timer <= 0:
                self.state = 'ON'

        self.power = self.P_rate if self.state in ['ON', 'ON_LOCKED'] else 0.0
        return self.state, self.power

class VirtualCluster:
    def __init__(self, num_units, P_rate, lock_time, execution_period, T_o):
        self.units = [
            SemiMarkovAC(
                P_rate,
                lock_time,
                execution_period,
                ETPModel(
                    C_a=1e3, C_m=2e3,
                    H_m=50.0,
                    U_a=100.0,
                    COP=3.5,
                    T_o=T_o
                )
            )
            for _ in range(num_units)
        ]
        self.execution_period = execution_period
        self.lock_time = lock_time
        self.P_rate = P_rate

    def step(self, target_response):
        expected_powers = np.array([unit.compute_expected_power() for unit in self.units])
        average_expected_power = np.mean(expected_powers)

        # 推导集群开启概率
        p_ratio = target_response / self.P_rate
        # 如果响应功率较大，那么u_on采用公式Rand(0.5Δt/t_lock, 1.5Δt/t_lock)得到，
        # 反之则u_off通过公式得到，然后另一个迁移概率通过稳态功率公式计算得到。
        if p_ratio > 0.5:
            # u_0 先由多样性约束随机采样
            u_0 = np.random.uniform(0.5*self.execution_period/self.lock_time, 1.5*self.execution_period/self.lock_time)
            T_1 = self.execution_period / u_0  #
            # 由如下空调稳态期望功率可得
            # target_response = self.P_rate*(T_1+self.lock_time)/(T_1+self.lock_time*2+self.execution_period/u_1)

            u_1 = self.execution_period / (target_response / self.P_rate / (T_1+self.lock_time) - (T_1 + self.lock_time*2))

        else:
            # 类似可得
            u_1 = np.random.uniform(0.5*self.execution_period/self.lock_time, 1.5*self.execution_period/self.lock_time)
            T_2 = self.execution_period / u_0
            u_0 = self.execution_period / (target_response / self.P_rate / (T_2+self.lock_time) - (T_2 + self.lock_time*2))

        # 执行控制
        states = []
        powers = []
        for unit, u_0 in zip(self.units, u_0_list):
            state, power = unit.update_state(u_0, u_1)
            states.append(state)
            powers.append(power)
            unit.etp.step(Q_ac=power * unit.etp.COP, Q_i=0, Q_s=0, dt=self.execution_period)

        total_power = np.sum(powers)
        return states, powers, total_power